import torch
import torch.optim as optim
from torch.autograd import Variable
import torch.nn as nn
import numpy as np
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


class simpleNet(nn.Module):
    # 定义单隐藏层的神经网络，我们预想这个网络应该有一个16个节点的隐含层，每个节点代表一个解析几何的平面方程
    # 我并不确定我这样写是否是正确的网络搭建形式，python我刚入门不到3周时间。。
    def __init__(self, in_dim, n_hidden, out_dim):
        super(simpleNet, self).__init__()
        #定义全连接隐含层，输入3，输出16，后面Ture是启用偏置
        self.hiddenlayer = nn.Linear(in_dim, n_hidden, True)
        #输出层输入16，输出2
        self.outlayer = nn.Linear(n_hidden, out_dim, False)
        #隐含层用ReLu激活
        net = [self.hiddenlayer, nn.ReLU()]
        net += [self.outlayer]
        self.model = nn.Sequential(*net)

    def forward(self, x):
        #softmax对输出结果处理，实际上是预测这个体素是否在模型内，softmax函数让两个输出节点的和为1
        #网络训练结束之后，我们将整个空间中的体素方块都通过这个网络做，输出结果其实就是代表了该体素方块是否是在模型内
        #这样我们就得到了整个数据空间的体素分布情况。。。
        output = self.model(x)
        #血的教训，不需要用softmax处理，交叉熵loss自带softmax，两个softmax会让梯度消失导致loss不下降
        'output = torch.softmax(output, dim=-1)'
        return output